#!/usr/bin/perl
#
# fastq-subsample.pl
#
# This script subsamples a fastq or pair of paired-end fastq files. 
# The sampling rate is passed as an argument to the script. 
# 
# Copyright (C) 2014 Shawn Driscoll, sdriscoll@salk.edu
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
#

use strict;
use Getopt::Std;

# -- check args
my $argc = scalar @ARGV;
my %opts = (
	"a" => "",
	"o" => "subsample"
);

# check options
my $r = getopt("o:a:", \%opts);
if(!$r) {
	die usage();
}

die usage() if $argc < 2;

# -- variables

my $rname;
my $rrand;
my $rate = shift @ARGV;
my $fq1 = shift @ARGV;
my $out1 = "";
my $out2 = "";
my $fq2 = "";
my $pe = 0;
my $fqline = 0;
if(scalar @ARGV) {
	$fq2 = shift @ARGV;
	$out2 = $fq1 . ".sub";
	$pe = 1;
}
my $gz = $fq1 =~ /\.gz$/ ? 1 : 0;
my $rename = $opts{"a"} ne "" ? 1 : 0;
my $i;
my $rcnt = 0;
my $r2 = "";

# output file name(s)
if($pe) {
	$out1 = $opts{"o"} . ".1.fq";
	$out2 = $opts{"o"} . ".2.fq";
}
else {
	$out1 = $opts{"o"} . ".fq";
}

# -- open file(s) and sample lines

if($gz) {
	open IN1, "gunzip -c $fq1 |";
	if($pe) {
		open IN2, "gunzip -c $fq2 |";
	}
}
else {
	open IN1, '<', $fq1;
	if($pe) {
		open IN2, '<', $fq2;
	}
}

open OUT1, '>', $out1;
if($pe) {
	open OUT2, '>', $out2;
}

print STDERR "# subsampling rate $rate\n";

$fqline = 0;
while(<IN1>) {
	$fqline++;
	if($pe) {
		$r2 = <IN2>;
	}

	if(rand() < $rate) {
		# print this read (or pair of reads)
		if($rename) {
			$rname = "@" . $opts{"a"} . ":$rcnt";
			print OUT1 $rname . "/1" . "\n";
			if($pe) {
				print OUT2 $rname . "/2" . "\n";
			}
		}
		else {
			print OUT1 $_;
			if($pe) {
				print OUT2 $r2;
			}
		}

		$i = 3;
		while($i--) {
			$_ = <IN1>;
			print OUT1 $_;
			if($pe) {
				$r2 = <IN2>;
				print OUT2 $r2;
			}
		}
		$rcnt++;		
	}
	else {
		# read 3 more lines and loop back around
		$i = 3;
		while($i--) {
			$_ = <IN1>;
			if($pe) {
				$_ = <IN2>;
			}
		}
	}

}

close IN1;
close OUT1;
if($pe) {
	close IN2;
	close OUT2;
}

print STDERR "# wrote $rcnt reads\n";

# -- subs

sub usage {
	return "
Usage: subsample-fq.pl [options] rate fq1 [matefq]

Options
  -a STR    Stub name to use to rename the reads
  -o STR    Stub name for output files (default: subsample)
";
}
